// npu_dotp_relu.sv
// 支持 AXI4-Lite 访问主存的完整 NPU 加速器
`timescale 1ns/1ps

module npu_dotp_relu import riscv_pkg::*; (
  input  logic        clk,
  input  logic        rst_n,

  // CSR 接口（来自 RISC-V 核心）
  input  logic        cs_i,
  input  logic        we_i,
  input  logic [7:0]  addr_i,       // 0x00 ~ 0xFF
  input  logic [31:0] wdata_i,
  output logic [31:0] rdata_o,
  output logic        ack_o,

  // AXI4-Lite Master 接口（连接到系统总线）
  // 地址：32-bit, 数据：32-bit
  output logic [31:0] axi_awaddr,
  output logic        axi_awvalid,
  input  logic        axi_awready,

  output logic [31:0] axi_wdata,
  output logic [3:0]  axi_wstrb,
  output logic        axi_wvalid,
  input  logic        axi_wready,

  input  logic [1:0]  axi_bresp,
  input  logic        axi_bvalid,
  output logic        axi_bready,

  output logic [31:0] axi_araddr,
  output logic        axi_arvalid,
  input  logic        axi_arready,

  input  logic [31:0] axi_rdata,
  input  logic [1:0]  axi_rresp,
  input  logic        axi_rvalid,
  output logic        axi_rready
);

  // ========================================================
  // 1. CSR 寄存器定义
  // ========================================================
  localparam logic [7:0]
    ADDR_CTRL   = 8'h00,  // [0]=start, [1]=done(ro), [31:16]=vec_len
    ADDR_STATUS = 8'h04,  // [0]=busy, [1]=done
    ADDR_SRC_A  = 8'h10,  // word addr
    ADDR_SRC_B  = 8'h14,
    ADDR_DST    = 8'h18,
    ADDR_RESULT = 8'h20;  // ro

  logic [31:0] reg_ctrl, reg_src_a, reg_src_b, reg_dst;
  logic [15:0] vec_len;
  logic        start_pulse;

  // 写寄存器
  always_ff @(posedge clk or negedge rst_n) begin
    if (!rst_n) begin
      reg_ctrl  <= 32'd0;
      reg_src_a <= 32'd0;
      reg_src_b <= 32'd0;
      reg_dst   <= 32'd0;
    end else if (cs_i && we_i) begin
      case (addr_i)
        ADDR_CTRL:   reg_ctrl  <= wdata_i;
        ADDR_SRC_A:  reg_src_a <= wdata_i;
        ADDR_SRC_B:  reg_src_b <= wdata_i;
        ADDR_DST:    reg_dst   <= wdata_i;
        default: ;
      endcase
    end
  end

  assign vec_len = reg_ctrl[31:16];
  // 检测 start 上升沿
  logic start_last;
  always_ff @(posedge clk) start_last <= reg_ctrl[0];
  assign start_pulse = reg_ctrl[0] && !start_last && cs_i && we_i && (addr_i == ADDR_CTRL);

  // ========================================================
  // 2. 状态机
  // ========================================================
  typedef enum logic [3:0] {
    IDLE,
    READ_A,      // 发 AR (A)
    WAIT_R_A,    // 等待 R (A)
    READ_B,      // 发 AR (B)
    WAIT_R_B,    // 等待 R (B)
    COMPUTE,
    WRITE_DST,   // 发 AW + W
    WAIT_BRESP,
    DONE
  } state_t;

  state_t cstate, nstate;
  always_ff @(posedge clk or negedge rst_n) begin
    if (!rst_n) cstate <= IDLE;
    else       cstate <= nstate;
  end

  // ========================================================
  // 3. 计数器 & 索引
  // ========================================================
  localparam int SIMD_W = 8;
  logic [15:0] idx;                    // 当前处理的向量索引
  logic [15:0] word_offset_a, word_offset_b;
  logic        last_chunk;

  assign word_offset_a = idx;
  assign word_offset_b = idx;
  assign last_chunk    = (idx + SIMD_W >= vec_len);

  always_ff @(posedge clk) begin
    if (cstate == IDLE && start_pulse)
      idx <= 0;
    else if (cstate == COMPUTE && !last_chunk)
      idx <= idx + SIMD_W;
  end

  // ========================================================
  // 4. AXI Read Channel (A and B)
  // ========================================================
  logic [31:0] mem_data_a, mem_data_b;
  logic        rd_a_valid, rd_b_valid;

  // 读 A
  assign axi_araddr  = (cstate == READ_A)  ? (reg_src_a + {word_offset_a, 2'b00}) :
                       (cstate == READ_B)  ? (reg_src_b + {word_offset_b, 2'b00}) : 32'd0;
  assign axi_arvalid = (cstate == READ_A || cstate == READ_B);

  assign axi_rready  = 1'b1;  // 始终准备好接收

  always_ff @(posedge clk) begin
    if (axi_arvalid && axi_arready) begin
      if (cstate == READ_A)  rd_a_valid <= 1'b0;
      if (cstate == READ_B)  rd_b_valid <= 1'b0;
    end
    if (axi_rvalid) begin
      if (cstate == WAIT_R_A) begin
        mem_data_a <= axi_rdata;
        rd_a_valid <= 1'b1;
      end
      if (cstate == WAIT_R_B) begin
        mem_data_b <= axi_rdata;
        rd_b_valid <= 1'b1;
      end
    end
  end

  // ========================================================
  // 5. SIMD 数据提取
  // ========================================================
  logic signed [15:0] vec_a [0:SIMD_W-1];
  logic signed [15:0] vec_b [0:SIMD_W-1];
  logic [7:0] valid_lanes;

  always_comb begin
    for (int i = 0; i < SIMD_W; i++) begin
      vec_a[i] = (i < (vec_len - idx) && idx + i < vec_len) ?
                  $signed(mem_data_a[16*i +: 16]) : 16'sd0;
      vec_b[i] = (i < (vec_len - idx) && idx + i < vec_len) ?
                  $signed(mem_data_b[16*i +: 16]) : 16'sd0;
    end
    valid_lanes = (vec_len - idx > SIMD_W) ? 8'd8 : (vec_len - idx);
  end

  // ========================================================
  // 6. 乘加树 + 累加器
  // ========================================================
  logic signed [31:0] mul [0:SIMD_W-1];
  logic signed [31:0] partial_sum;
  logic signed [63:0] acc_reg;  // 防止溢出

  always_comb begin
    partial_sum = 0;
    for (int i = 0; i < SIMD_W; i++) begin
      mul[i] = vec_a[i] * vec_b[i];
      if (i < valid_lanes)
        partial_sum = partial_sum + mul[i];
    end
  end

  always_ff @(posedge clk) begin
    if (cstate == IDLE && start_pulse)
      acc_reg <= 64'd0;
    else if (cstate == COMPUTE)
      acc_reg <= acc_reg + partial_sum;
  end

  logic [31:0] result;
  assign result = acc_reg[31] ? 32'd0 : acc_reg[31:0];  // ReLU

  // ========================================================
  // 7. AXI Write Channel (写回 DST)
  // ========================================================
  assign axi_awaddr  = reg_dst;
  assign axi_awvalid = (cstate == WRITE_DST);
  assign axi_wdata   = result;
  assign axi_wstrb   = 4'b1111;
  assign axi_wvalid  = (cstate == WRITE_DST);
  assign axi_bready  = (cstate == WAIT_BRESP);

  // ========================================================
  // 8. 状态转移逻辑
  // ========================================================
  always_comb begin
    nstate = cstate;
    case (cstate)
      IDLE:       if (start_pulse) nstate = READ_A;

      READ_A:     if (axi_arvalid && axi_arready) nstate = WAIT_R_A;

      WAIT_R_A:   if (rd_a_valid) nstate = READ_B;

      READ_B:     if (axi_arvalid && axi_arready) nstate = WAIT_R_B;

      WAIT_R_B:   if (rd_b_valid) nstate = COMPUTE;

      COMPUTE:    nstate = last_chunk ? WRITE_DST : READ_A;

      WRITE_DST:  if (axi_awvalid && axi_awready &&
                      axi_wvalid  && axi_wready)  nstate = WAIT_BRESP;

      WAIT_BRESP: if (axi_bvalid) nstate = DONE;

      DONE:       nstate = IDLE;
    endcase
  end

  // ========================================================
  // 9. 状态标志
  // ========================================================
  logic busy, done;
  assign busy = (cstate != IDLE);
  assign done = (cstate == DONE);

  // 写 CTRL[1] 置位 done
  always_ff @(posedge clk) begin
    if (cstate == DONE)
      reg_ctrl[1] <= 1'b1;
    else if (cs_i && we_i && addr_i == ADDR_CTRL)
      reg_ctrl[1] <= wdata_i[1];  // 软件可清零
  end

  // ========================================================
  // 10. CSR 读端口
  // ========================================================
  always_comb begin
    rdata_o = 32'd0;
    ack_o   = cs_i;
    if (cs_i && !we_i) begin
      case (addr_i)
        ADDR_CTRL:   rdata_o = reg_ctrl;
        ADDR_STATUS: rdata_o = {30'd0, busy, done};
        ADDR_SRC_A:  rdata_o = reg_src_a;
        ADDR_SRC_B:  rdata_o = reg_src_b;
        ADDR_DST:    rdata_o = reg_dst;
        ADDR_RESULT: rdata_o = result;
        default:     rdata_o = 32'd0;
      endcase
    end
  end

  // 清零 start 位（软件写1后自动清零）
  always_ff @(posedge clk) begin
    if (start_pulse)
      reg_ctrl[0] <= 1'b0;
  end

endmodule